#!/bin/bash
# -------------------------------------------------------------------------------
# Filename:     start_web_demo.sh
# Update:       2025/04/16
# Description:  一键启动Web Demo服务，支持自定义参数配置
# Example:      ./start_web_demo.sh [--model MODEL_NAME] [--port 7860] [--ckpt /model/path]
#   默认启动：   ./start_web_demo.sh
#   参数启动：   ./start_web_demo.sh --model qwen2.5_omni  #用于后缀中的【模型名称】与【目录名称】有差别时使用。
#   参数启动：   ./start_web_demo.sh --port 8888 --ckpt /data/models/llm/models/Qwen2.5-Omni-7B
#   查看帮助：   ./start_web_demo.sh --help
# Depends:      离线Docker镜像, Python, Torch2.5环境, run_web_demo.sh, web_demo.py
# -------------------------------------------------------------------------------

set -eo pipefail

# -------------------------- 颜色定义 --------------------------
RED='\033[31m'
GREEN='\033[32m'
YELLOW='\033[33m'
CYAN='\033[36m'
NC='\033[0m'

# -------------------------- 基础配置 --------------------------
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
ROOT_DIR=$(dirname "${SCRIPT_DIR}")
SCRIPT_FILE="run-container-dev.sh"
RUN_SCRIPT="${ROOT_DIR}/${SCRIPT_FILE}"
CONTAINER_SCRIPT="/home/share/pytorch2.5/qwen2.5-omni/tools/run_web_demo.sh"
LOG_FILE="web_demo_start.log"
ENV_FILE="${ROOT_DIR}/env.sh"

# -------------------------- 默认参数 --------------------------
DEFAULT_MODEL_NAME="qwen2.5_omni"
DEFAULT_PORT=7860
DEFAULT_CKPT="/data/models/llm/models/Qwen2.5-Omni-7B"
DEFAULT_DEVICES="0,1,2,3"

# -------------------------- 函数定义 --------------------------
error() {
  echo -e "${RED}[ERROR]${NC} $1" >&2
  exit 1
}

info() {
  echo -e "${CYAN}[INFO]${NC} $1"
}

success() {
  echo -e "${GREEN}[SUCCESS]${NC} $1"
}

check_dependency() {
  if ! command -v docker &>/dev/null; then
    error "Docker未安装! 请先安装Docker环境"
  fi

  if [ ! -f "${RUN_SCRIPT}" ]; then
    error "缺失核心脚本: ${RUN_SCRIPT}"
  fi

  if [ ! -f "${ENV_FILE}" ]; then
    error "环境配置文件 ${ENV_FILE} 不存在!"
  fi
}

update_env() {
  local model_name="$1"  # 接收外部传入的模型名称

  # 如果未指定模型名称，则从根目录名称提取
  if [ -z "${model_name}" ]; then
    model_name=$(basename "${SCRIPT_DIR}")
    info "自动提取模型名称: ${model_name}"
  fi

  info "正在更新环境配置: TestModel=-${model_name}"
  pushd "${ROOT_DIR}" > /dev/null

  # 注释掉原始的 TestModel="" 行
  sed -i 's/^TestModel=""$/#TestModel=""/' "${ENV_FILE}" || error "注释TestModel失败"
  
  # 取消注释包含目标模型的行 (如果存在)
  sed -i "s/^#\(TestModel=\"-${model_name}\"\)/\1/" "${ENV_FILE}" || error "取消注释失败"
  
  # 直接更新现有值 (如果未注释版本存在)
  sed -i "s/^TestModel=.*$/TestModel=\"-${model_name}\"/" "${ENV_FILE}" || error "更新TestModel失败"

  popd > /dev/null
}

run_in_background() {
  info "正在后台启动容器服务..."
  
  #pushd ${ROOT_DIR} > /dev/null
  # 使用nohup后台运行并捕获输出
  nohup ${RUN_SCRIPT} \
      --cmd "${CONTAINER_SCRIPT} --port ${PORT} --ckpt ${CKPT} --devices ${DEVICES}" \
      --non-interactive >> "${LOG_FILE}" 2>&1 &

  # 记录后台进程PID
  BG_PID=$!
  info "后台进程PID: ${BG_PID}"
  #popd > /dev/null

  # 等待初始启动日志生成
  #sleep 2
  #tail -n 20 "${LOG_FILE}"
}

usage() {
  cat <<EOF

${GREEN}Usage:${NC} $0 [options]

${YELLOW}Options:${NC}
  --model <name>      指定模型名称 (优先级高于目录名，默认: 根目录名)
  --port <number>     Web服务端口 (默认: ${DEFAULT_PORT})
  --ckpt <path>       模型checkpoint路径 (默认: ${DEFAULT_CKPT})
  --devices <list>    MLU设备列表 (默认: ${DEFAULT_DEVICES})
  -h, --help          显示帮助信息

${YELLOW}示例:${NC}
  $0 --model my_model --port 8888 --ckpt /custom/model/path
  $0 --port 8888 --ckpt /custom/model/path --devices 0,1
  $0 --devices 0,1    # 使用目录名作为模型名称
EOF
  exit 0
}

# -------------------------- 参数解析 --------------------------
MODEL_NAME=""
PORT="${DEFAULT_PORT}"
CKPT="${DEFAULT_CKPT}"
DEVICES="${DEFAULT_DEVICES}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --model) MODEL_NAME="$2"; shift 2 ;;
    --port) PORT="$2"; shift 2 ;;
    --ckpt) CKPT="$2"; shift 2 ;;
    --devices) DEVICES="$2"; shift 2 ;;
    -h|--help) usage ;;
    *) error "未知参数: $1" ;;
  esac
done

# -------------------------- 执行流程 --------------------------
{
  # 预检查阶段
  info "正在检查系统依赖..."
  check_dependency
  
  # 动态更新环境变量
  update_env "${MODEL_NAME}"
  #update_env "qwen2.5_omni"

  # 后台启动服务
  run_in_background

  success "服务已在后台启动，模型加载完成后即可访问!"
  echo -e "\n${YELLOW}访问地址:${NC} http://localhost:${PORT}"
  echo -e "${YELLOW}实时日志:${NC} tail -f ${LOG_FILE}"
  echo -e "${YELLOW}进程管理:${NC}"
  echo -e "  查看状态: ps -p ${BG_PID}"
  echo -e "  终止服务: kill ${BG_PID}"
} | tee -a "${LOG_FILE}" 2>&1 
